"""
    构建huffman tree
"""
class HuffmanNode:
    def __init__(self, word_id, frequency):
        """
        初始化huffman树结点的数据结构
        :param word_id:
        :param frequency:
        """
        # 单词id
        self.word_id = word_id
        # 单词对应的频率
        self.frequency = frequency
        # 左孩子
        self.left_child = None
        # 右孩子
        self.right_child = None
        # 父节点
        self.father = None
        # huffman code
        self.Huffman_code = []
        # 路径
        self.path = []


class HuffmanTree:
    """
    初始化huffman树结构
    """
    def __init__(self, wordid_frequency_dict):
        """
        参数初始化
        :param wordid_frequency_dict: 词频率字典，key是单词id，值为频率
        """
        # 单词总数
        self.word_count = len(wordid_frequency_dict)
        # key是单词id，value为code，表示是怎么走的，往左走为0，往右走为1
        self.wordid_code = dict()
        # key是单词id，value为path,表示经过哪些结点
        self.wordid_path = dict()
        # 根节点
        self.root = None
        # 初始化未合并节点
        unmerge_node_list = [HuffmanNode(wordid, frequency) for wordid, frequency in wordid_frequency_dict.items()]
        # 初始化huffman树
        self.huffman = [HuffmanNode(wordid, frequency) for wordid, frequency in wordid_frequency_dict.items()]
        print("Building huffman tree...")
        self.build_tree(unmerge_node_list)
        print("Building tree finished")
        print("Generating huffman path...")
        # 生成huffman path
        self.generate_huffman_code_and_path()
        print("Generating huffman path finished")

    def merge_node(self, node1, node2):
        """
        合并两个节点
        :param node1:节点1
        :param node2:节点2
        :return:
        """
        # 将两个节点频度相加得到父亲节点的频度
        sum_frequency = node1.frequency + node2.frequency
        # 父亲节点的id
        mid_node_id = len(self.huffman)
        # 构建父亲节点
        father_node = HuffmanNode(mid_node_id, sum_frequency)
        # 孩子节点中节点频度大的为左孩子，节点频度小的为右孩子
        if node1.frequency >= node2.frequency:
            father_node.left_child = node1
            father_node.right_child = node2
        else:
            father_node.left_child = node2
            father_node.right_child = node1
        # huffman树添加父亲节点
        self.huffman.append(father_node)
        # 返回父亲节点
        return father_node

    def build_tree(self, node_list):
        """
        构造根节点
        :param node_list:未合并节点列表
        """
        # 当未合并节点列表长度大于1时，说明继续合并
        while len(node_list) > 1:
            i1 = 0  # 频度最小的节点
            i2 = 1  # 频度第二小的节点
            # 确保i1的频度小于i2的频度
            if node_list[i2].frequency < node_list[i1].frequency:
                [i1, i2] = [i2, i1]
            # 遍历后续节点，调整位置，使的i1的频度小于i2的频度并且i1、i2是最小的两个频度
            for i in range(2, len(node_list)):
                if node_list[i].frequency < node_list[i2].frequency:
                    i2 = i
                    if node_list[i2].frequency < node_list[i1].frequency:
                        [i1, i2] = [i2, i1]
            # 合并最小的两个节点
            father_node = self.merge_node(node_list[i1], node_list[i2])
            # 先删除大序号的，再删除小序号的，防止序号改变
            if i1 < i2:
                node_list.pop(i2)
                node_list.pop(i1)
            elif i1 > i2:
                node_list.pop(i1)
                node_list.pop(i2)
            else:
                raise RuntimeError('i1 should not be equal to i2')
            # 插入新节点
            node_list.insert(0, father_node)
        # 当循环结束时，node_list中只有一个节点即为根节点
        self.root = node_list[0]

    def generate_huffman_code_and_path(self):
        """
        构建code与path
        """
        # 从根节点开始遍历
        stack = [self.root]
        # 如果栈非空，则继续遍历
        while len(stack) > 0:
            # 弹出栈的最后一个节点
            node = stack.pop()
            # 顺着左子树走
            while node.left_child or node.right_child:
                code = node.Huffman_code
                path = node.path
                # 构造左孩子与右孩子的code与path
                node.left_child.Huffman_code = code + [1]
                node.right_child.Huffman_code = code + [0]
                node.left_child.path = path + [node.word_id]
                node.right_child.path = path + [node.word_id]
                # 把没走过的右子树加入栈
                stack.append(node.right_child)
                # 继续遍历左孩子
                node = node.left_child
            # 当循环结束时，均为叶子节点，即单词节点
            # 叶子节点的单词id
            word_id = node.word_id
            # 叶子节点的code
            word_code = node.Huffman_code
            # 叶子节点的路径
            word_path = node.path
            self.huffman[word_id].Huffman_code = word_code
            self.huffman[word_id].path = word_path
            # 把节点计算得到的霍夫曼码、路径  写入词典的数值中
            self.wordid_code[word_id] = word_code
            self.wordid_path[word_id] = word_path


    def get_all_pos_and_neg_path(self):
        """
        获取所有词的正向节点id和负向节点id数组
        :return:
        """
        # 所有词的正向路径数组
        positive = []
        # 所有词的负向路径数组
        negative = []
        for word_id in range(self.word_count):
            # 存放一个词 路径中的正向节点id
            pos_id = []
            # 存放一个词 路径中的负向节点id
            neg_id = []
            for i, code in enumerate(self.huffman[word_id].Huffman_code):
                # 如果code=1，则为正样本
                if code == 1:
                    pos_id.append(self.huffman[word_id].path[i])
                # 否则，为负样本
                else:
                    neg_id.append(self.huffman[word_id].path[i])
            positive.append(pos_id)
            negative.append(neg_id)
        # 返回所有词的正向节点id和负向节点id数组
        return positive, negative


if __name__ == "__main__":
    word_frequency = {0: 4, 1: 6, 2: 3, 3: 2, 4: 2}
    print(word_frequency)
    # 构建huffman树
    tree = HuffmanTree(word_frequency)
    print(tree.wordid_code)
    print(tree.wordid_path)
    for i in range(len(word_frequency)):
        print(tree.huffman[i].path)
    # 获取所有词的正向节点id和负向节点id数组
    print(tree.get_all_pos_and_neg_path())

